Skip to content

Bump Nx to 0.13 - #358

Merged
josevalim merged 17 commits into
elixir-nx:mainfrom
RicardoSantos-99:bump-nx-0-13
Aug 23, 2026
Merged

Bump Nx to 0.13#358
josevalim merged 17 commits into
elixir-nx:mainfrom
RicardoSantos-99:bump-nx-0-13

Conversation

@RicardoSantos-99

@RicardoSantos-99 RicardoSantos-99 commented Aug 22, 2026

Copy link
Copy Markdown
Contributor

Bumps Nx and EXLA to 0.13

Draft while the doctests and test fixtures that shifted with the newer numerics get verified one by one against a reference implementation.

Nx and EXLA 0.13 declare elixir ~> 1.17, so the project requirement and the
older leg of the CI matrix move up with them.
Nx 0.13 prints f32 with the shortest representation that round trips, so
7.071067810058594 became 7.071068. The values are the same, only the printed
form changed.

Applies to the doctests where the numbers are identical as f32. The ones whose
values actually moved are left for a separate commit, since those need to be
justified rather than accepted.
Same formatting-only change as the previous commit, for doctests that
diverged after the branch was rebased onto a main that had precision
fixes landed since (elixir-nx#354). Values are unchanged, only the number of
significant digits Nx prints for f32.
Verified against sklearn on the iris dataset (fit, incremental_fit,
partial_fit, transform) and against a numpy SVD reference (fit_transform,
n=6). The n=6 case is the clearest evidence this is a real precision fix,
not just reformatting: the old second-component value was off by ~0.045
from the numpy reference, the new one matches to ~1e-6. Small-sample PCA
is exactly the ill-conditioned regime the batched eigh bug (found while
working on QDA, elixir-nx#349) would be expected to hurt most.
Verified against a numpy reference matching the module's own centered
kernel eigendecomposition and projection formulas. The new values land
within 1e-7 to 1e-8 of the true eigensolution across all three doctests
(fit, transform, fit_transform); the old ones were off by 1e-4 to 1e-5,
consistently about three orders of magnitude worse. Same eigh precision
fix as PCA.
Verified against the exact (non-randomized) SVD of each test's own
generated matrix, which the earlier bug hunt already established as
the right reference here (Scholar's randomized SVD is faithful to
whatever Nx computes; the question is whether Nx computes it well).

New values match the true SVD to ~1e-5 to 1e-7 across all 7 fixtures
(2 doctests, 5 tests). Old values were off by 1e-3 to 2.7e-2, several
past the tests' own 1e-3 tolerance, concentrated on the second
singular vector, the same pattern as PCA and KernelPCA.
Matches the analytical solution exactly now (coef=[-0.5,-0.7],
intercept=5.9 for this exactly-determined 3x3 system, verified before
the rebase too). Not touching the n=1 sample_weights test: its old
fixture (coefficient 0.0) is what sklearn gives for that degenerate
input, and the new computed value (1.9999996) is wrong, not improved.
Root cause: Nx.weighted_mean's centering leaves a ~3e-8 residual under
0.13 instead of the exact 0 the old Nx version happened to round to
for this input, and that residual is above Nx.LinAlg.pinv's default
1e-10 zero threshold, so pinv inverts real noise instead of treating
it as zero. Pre-existing fragility in the centering step, exposed by
the bump, not caused by it. Needs its own fix, not a fixture update.
Verified against sklearn.linear_model.Ridge(alpha=1.0, solver='svd')
on the same data. New values match to ~1e-7 to 1e-8; old ones were off
by 1e-3 to 1.75e-3.
Same underlying LinearRegression solver on the expanded feature matrix.
degree:1 matches the analytical solution exactly, same as
LinearRegression. degree:2 verified against a numpy minimum-norm pinv
reference on the same expanded features (3 samples, 5 features,
underdetermined). New matches to ~1e-7/1e-8; old was off by up to 3.7e-5
on coefficients and 3e-4 on the intercept.
Different verification than the other modules: t-SNE is a non-convex
iterative optimization with no unique correct answer, so there is no
ground truth to check the new values against. Tiny floating-point
differences across many gradient steps chaotically diverge into a
different (but not necessarily worse) local optimum.

Verified instead by: the doctest's perfectly symmetric input (4
equidistant collinear points) still produces the exact same
antisymmetric structure under 0.13, with only ~1e-6 relative drift, so
the underlying computation is sound. For the 7 test fixtures, which do
diverge substantially in absolute numbers, checked embedding quality
directly (k=3 neighbor overlap and rank correlation with the original
pairwise distances) on one case: comparable to the old fixture, not
degraded. No NaN or Inf in any of the 7 new embeddings.
Small, well-conditioned tridiagonal solve for the control points.
Values shifted by ~1e-6 to 2e-3, ordinary f32 rounding-level drift,
no ill-conditioning concern like PLSSVD's near-degenerate case.
Pure formatting, same values, shorter f32 representation. The
earlier automated formatting pass missed these (single-occurrence
scalars/vectors further down the same doctests, not caught by that
pass's token matching).
Pure formatting, same value, shorter f32 representation, same class
of miss as the LedoitWolf fix.
Pure formatting, same values, shorter f32 representation. Same class
of miss as LedoitWolf/ShrunkCovariance: single-occurrence values
further down a doctest with a repeated literal elsewhere in the file
(0.6000000238418579 also appears as a Nx.f32() input argument
elsewhere, left untouched).
… input

With a single sample (or samples that coincide after weighting),
centering is mathematically exact zero, but a - weighted_mean(a) is
not guaranteed to round-trip to exact 0.0 in floating point the way
a - a is. Nx.LinAlg.pinv has no way to tell that residual apart from
a real near-zero singular value, since it only sees the already-tiny
centered scale, not the original data's scale. Under Nx 0.9.2 this
particular case happened to round to exact 0 and the bug was never
visible; under 0.13 it lands a few ULPs off, clears pinv's default
1e-10 zero threshold, and gets inverted into a large spurious
coefficient (a division by near-zero) instead of the correct 0.

Snap the centered design matrix to zero using a tolerance relative to
its own pre-centering scale, the same idea numpy/scipy use for
lstsq's rank tolerance, before it reaches pinv. Verified this does not
touch well-conditioned fits (analytical 3x3 system, a random 20x3
regression) or a genuinely tiny real slope (coefficient ~0.1 from
x values at 1e-4 scale, unaffected since the tolerance is relative to
x's own scale, not some larger unrelated one). Not Nx-bump specific:
this was always a latent gap in centering's exactness assumption, the
version bump only stopped masking it by accident.
Nx.LinAlg.svd's full_matrices?: false path forms the Gram matrix CᵀC
to avoid the larger QR/Halley iteration, which squares the condition
number, a trade-off the function's own source comments call out
explicitly. PLSSVD's cross-covariance matrix routinely has singular
values far apart (6.8 and 0.014 on the module's own doctest data, a
~470:1 ratio), squaring that gap is enough for the smaller singular
vector to lose real precision: measured directly, x_weights' second
column came back with norm 0.9967 instead of 1.0, not a rounding
difference. Reproduces identically on the binary backend, so it is
Nx.LinAlg.SVD's algorithm, not an EXLA-specific issue.

full_matrices?: true takes the QDWH-based path instead of forming the
Gram matrix, at the cost of more work. Verified against a numpy
reference on both the module's doctest data and the test suite's
larger dataset: new values match to ~1e-6 to 1e-7 (down from old
errors of 1e-4 to 1.5e-3), and the resulting x_weights and y_weights
are now genuinely orthonormal (column norm 0.9999999, cross column
dot product -2.9e-8). Not Nx-bump specific: the same defect reproduces
identically under Nx 0.9.2 on the isolated matrix, it just never
surfaced before because nothing in Scholar exercised this precise
combination of full_matrices?: false and a poorly separated cross-
covariance before this verification pass.
Same root cause as PLSSVD: Nx.LinAlg.svd's full_matrices?: false path
forms the Gram matrix AᵀA, squaring the condition number, at the cost
of precision in the smaller singular vectors. Confirmed real, not just
imprecise: on a matrix with one dominant direction and three ~1/500
scale directions, the alpha=1.0 regularization was assumed to damp any
such error away, but it does not fully. The :svd solver disagreed with
Scholar's own :cholesky solver (which never touches svd) on the three
small coefficients, one of them with the wrong sign. Cross-checked
against sklearn.linear_model.Ridge with both solvers, which agree with
each other and with Scholar's :cholesky path to ~1e-7; Scholar's :svd
path was off by up to an order of magnitude before this fix, matches
to ~1e-5/1e-6 after.

full_matrices?: true needs u and vt sliced back down to the compact
singular value count afterward, since it returns square U/V instead
of the economy shape the rest of solve_svd assumes.

Found via a systematic sweep (14 shapes x 10 condition ratios x 8
seeds) after a first pass of hand-picked test cases missed it; the
same sweep found zero reproductions through the real PCA.fit despite
finding the underlying defect in raw Nx.LinAlg.svd calls at PCA-like
shapes, so PCA was left alone.
@josevalim

Copy link
Copy Markdown
Contributor

Draft while the doctests and test fixtures that shifted with the newer numerics get verified one by one against a reference implementation.

FWIW, you only need to do so for the values that actually changed, as most changes are pretty printing related.

@RicardoSantos-99
RicardoSantos-99 marked this pull request as ready for review August 23, 2026 11:28
@RicardoSantos-99

Copy link
Copy Markdown
Contributor Author

Draft while the doctests and test fixtures that shifted with the newer numerics get verified one by one against a reference implementation.

FWIW, you only need to do so for the values that actually changed, as most changes are pretty printing related.

Yeah, that’s what I did for most of them. A handful looked a bit off though, so I dug into those specifically and ended up finding a few real bugs, unrelated to the bump itself. They were just numerical issues that got exposed by the newer SVD path. I fixed those too (LinearRegression, PLSSVD, RidgeRegression) and verified the results against sklearn and Scholar’s own solvers. The PR is out of draft now.

@RicardoSantos-99

Copy link
Copy Markdown
Contributor Author

While reviewing some of these modules, I found a few other bugs as well, but I didn’t include them in this PR since they’re unrelated to the bump. I still need to investigate them a bit further, and I’ll either open a separate PR or create an issue for them.

@josevalim
josevalim merged commit 71c7c80 into elixir-nx:main Aug 23, 2026
2 checks passed
@josevalim

Copy link
Copy Markdown
Contributor

💚 💙 💜 💛 ❤️

@RicardoSantos-99
RicardoSantos-99 deleted the bump-nx-0-13 branch August 23, 2026 18:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants